import os, sys
import numpy as np
import ncnn
import cv2
import ncnn as pyncnn
from pathlib import Path
from glob import glob
import time

from inference import Inference
from utilities import *
import ops

if __name__ == "__main__":
    # inf = Inference("weights/2in1_yolo11n_ncnn_model")
    # inf = Inference(f"{os.getcwd()}/weights/2in1_infrared_ncnn_model")
    inf = Inference(f"{os.getcwd()}/weights/yolo11n_ncnn_model")
    names = inf.names

    # files = glob("D:/datasets/eye_2in1/images/*.jpg")
    files = glob("images/*.jpg")
    for file in files:
        try:
            origin = cv2.imread(file)
            # origin = origin[160 : 160 + 320, 180 : 180 + 320]
        except Exception as e:
            print(f"Error reading image {file}: {e}")

        base_name = os.path.basename(file)

        if origin is None:
            continue

        # if origin is None:
        #     origin = cv2.imread(f"ncnn_inference/{file}")

        profiler = ops.Profile()
        with profiler:
            result = inf.inference(origin)

        print("run time:")
        print(profiler.dt)
        print("")

        im_show = origin
        for i, data in enumerate(result.data):
            cls, conf, bbox = result.cls[i], result.conf[i], result.data[i]
            # im_show = cv2.rectangle(origin, bbox, cls, conf, names)
            label = names[int(cls)]
            color = colors(int(cls), True)
            im_show = box_label(origin, bbox, label, conf, color)
        cv2.putText(im_show, f"{base_name}", (10, 30), cv2.FONT_HERSHEY_COMPLEX, 1, (50, 50, 220), 1)

        cv2.imshow("detect", im_show)
        cv2.waitKey(0)
